import time
import matplotlib.pyplot as plt
import torch
import numpy as np
import csv
import pandas as pd
from torch.utils.data import DataLoader, Dataset
import torch.nn as nn
from torch import optim
from sklearn.feature_selection import SelectKBest, chi2
from sklearn.decomposition import PCA
from sklearn.preprocessing import StandardScaler
import os
os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

# ================== 特征选择模块 ==================
def get_feature_importance(feature_data, label_data, k=4, column=None):
    """SelectKBest特征选择（基于卡方检验）"""
    model = SelectKBest(chi2, k=k)
    feature_data = np.array(feature_data, dtype=np.float64)
    X_new = model.fit_transform(feature_data, label_data)  # 自动选择k个最佳特征
    indices = np.argsort(model.scores_)[::-1][:k]  # 按得分倒序排列取前k
    
    if column:  # 打印特征名称
        k_best = [column[i + 1] for i in indices]
        print(f'Top {k} features:', k_best)
    return X_new, indices

def pca_reduction(feature_data, n_components=4):
    """PCA特征降维（处理高维共线性问题）"""
    pca = PCA(n_components=n_components)
    X_pca = pca.fit_transform(feature_data)  # 自动计算主成分
    print(f'主成分方差解释率: {pca.explained_variance_ratio_}')  # 显示各成分贡献度
    return X_pca, pca

# ================== 数据加载类 ==================
class CovidDataset(Dataset):
    def __init__(self, file_path, mode="train", feature_method='selectkbest', feature_dim=6):
        """初始化数据集（支持多种特征处理方法）
        - mode: train/val/test 模式
        - feature_method: selectkbest/pca 特征选择方法
        - feature_dim: 最终特征维度
        """
        with open(file_path, "r") as f:
            ori_data = list(csv.reader(f))
            column = ori_data[0]  # 保留列名用于特征解释
            csv_data = np.array(ori_data[1:])[:, 1:].astype(float)  # 去掉首列ID

        # 分离特征和标签（最后一列为标签）
        feature = np.array(ori_data[1:])[:, 1:-1]  # 特征列
        label_data = np.array(ori_data[1:])[:, -1]  # 标签列（test模式无）

        # ===== 特征处理路由 =====
        if feature_method == 'selectkbest':
            # 选择卡方检验得分最高的k个特征
            _, col_indices = get_feature_importance(feature, label_data, feature_dim, column)
            self.features = feature[:, col_indices]
        elif feature_method == 'pca':
            # PCA降维到指定维度
            self.features, self.pca_model = pca_reduction(feature, feature_dim)
        else:  # 不使用特征选择
            self.features = feature[:, :feature_dim]

        # ===== 数据标准化 ===== 
        # 每个特征单独标准化（均值0，方差1）
        self.features = StandardScaler().fit_transform(self.features)

        # ===== 数据划分 =====
        # 训练集取80%（5取4），验证集20%（逢5取1）
        if mode == "train":
            indices = [i for i in range(len(csv_data)) if i % 5 != 0]
        elif mode == "val":
            indices = [i for i in range(len(csv_data)) if i % 5 == 0]
        else:  # test模式使用全部数据
            indices = [i for i in range(len(csv_data))]

        # 转换为PyTorch张量
        self.features = torch.tensor(self.features[indices]).float()
        self.labels = torch.tensor(csv_data[indices, -1]).float() if mode != "test" else None

    def __getitem__(self, idx):
        """获取单个样本（测试集无标签）"""
        if self.labels is not None:
            return self.features[idx], self.labels[idx]
        return self.features[idx]

    def __len__(self):
        return len(self.features)

# ================== 模型定义 ==================
class CovidModel(nn.Module):
    """双层全连接网络结构
    - 输入维度: input_dim（根据特征选择结果动态调整）
    - 隐藏层: 64个神经元
    - 输出: 单值回归预测
    """
    def __init__(self, input_dim):
        super(CovidModel, self).__init__()
        self.layers = nn.Sequential(
            nn.Linear(input_dim, 64),  # 全连接层1
            nn.ReLU(),                 # 激活函数
            nn.Linear(64, 1)           # 输出层
        )

    def forward(self, x):
        return self.layers(x).squeeze()  # 去除多余维度

# ================== 正则化损失函数 ==================
def regularized_loss(pred, target, model, alpha=0.00075):
    """带L2正则化的MSE损失函数
    - alpha: 正则化强度系数（通过交叉验证确定）
    """
    mse = nn.MSELoss()(pred, target)
    l2_reg = torch.tensor(0.)
    for param in model.parameters():  # 累加所有参数的L2范数
        l2_reg += torch.norm(param)
    return mse + alpha * l2_reg  # 总损失 = MSE + λ*L2

# ================== 训练验证流程 ==================
def train_compare(feature_methods=['selectkbest', 'pca']):
    """对比不同特征选择方法的训练效果"""
    results = {}
    for method in feature_methods:
        print(f"\n=== 正在训练: {method.upper()} 方法 ===")
        
        # 数据加载（动态初始化）
        train_set = CovidDataset('covid.train.csv', 'train', feature_method=method)
        val_set = CovidDataset('covid.train.csv', 'val', feature_method=method)

        # 创建数据加载器（训练集打乱）
        train_loader = DataLoader(train_set, batch_size=32, shuffle=True)
        val_loader = DataLoader(val_set, batch_size=32)

        # 模型初始化（动态适配输入维度）
        model = CovidModel(input_dim=train_set.features.shape[1])
        optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

        # 训练循环（20个epoch）
        best_val_loss = float('inf')
        for epoch in range(20):
            model.train()  # 训练模式
            train_loss = 0
            for x, y in train_loader:
                optimizer.zero_grad()
                pred = model(x)
                loss = regularized_loss(pred, y, model)
                loss.backward()    # 反向传播
                optimizer.step()   # 参数更新
                train_loss += loss.item()

            # 验证阶段
            model.eval()  # 评估模式
            val_loss = 0
            with torch.no_grad():  # 禁用梯度计算
                for x, y in val_loader:
                    pred = model(x)
                    val_loss += regularized_loss(pred, y, model).item()

            # 记录最佳模型
            avg_val = val_loss / len(val_loader)
            if avg_val < best_val_loss:
                best_val_loss = avg_val
                torch.save(model.state_dict(), f'best_{method}.pth')  # 保存模型参数

            print(f'Epoch {epoch + 1}: 训练损失 {train_loss/len(train_loader):.4f} | 验证损失 {avg_val:.4f}')

        results[method] = best_val_loss

    # 可视化对比结果
    plt.figure(figsize=(10, 6))
    plt.bar(results.keys(), results.values(), color=['blue', 'orange'])
    plt.ylabel('验证集MSE')
    plt.title('特征选择方法对比')
    plt.savefig('method_comparison.png')
    plt.show()

# ================== 主程序 ==================
if __name__ == "__main__":
    # 运行对比实验
    train_compare()

    # PCA分析可视化（附加）
    full_data = np.array(pd.read_csv('covid.train.csv').iloc[:, 1:-1])
    pca = PCA().fit(StandardScaler().fit_transform(full_data))  # 全局PCA分析
    
    # 绘制累积方差解释曲线
    plt.figure()
    plt.plot(np.cumsum(pca.explained_variance_ratio_))
    plt.xlabel('主成分数量')
    plt.ylabel('累积方差解释率')
    plt.axhline(y=0.95, color='r', linestyle='--', label='95%阈值')
    plt.title('PCA方差解释分析')
    plt.savefig('pca_variance.png')